#!/usr/bin/env python3
"""
自动更新版本号的脚本
"""

import re
import sys
from pathlib import Path

def update_version(file_path, new_version):
    """更新文件中的版本号"""
    content = Path(file_path).read_text()
    
    # 匹配版本号模式
    if file_path.name == 'setup.py':
        pattern = r'(version\s*=\s*["\'])([^"\']+)(["\'])'
        replacement = rf'\g<1>{new_version}\g<3>'
    elif file_path.name == 'pyproject.toml':
        pattern = r'(version\s*=\s*["\'])([^"\']+)(["\'])'
        replacement = rf'\g<1>{new_version}\g<3>'
    else:
        raise ValueError(f"不支持的文件类型: {file_path}")
    
    updated_content = re.sub(pattern, replacement, content, count=1)
    Path(file_path).write_text(updated_content)
    print(f"已更新 {file_path} 中的版本号为 {new_version}")

def main():
    if len(sys.argv) != 2:
        print("用法: python update_version.py <新版本号>")
        sys.exit(1)
    
    new_version = sys.argv[1]
    
    # 定义需要更新的文件
    project_root = Path(__file__).parent
    setup_py = project_root / "setup.py"
    pyproject_toml = project_root / "pyproject.toml"
    
    # 更新版本号
    update_version(setup_py, new_version)
    update_version(pyproject_toml, new_version)
    
    print(f"版本号已成功更新为 {new_version}")

if __name__ == "__main__":
    main()
